# coding=utf-8
from keras.models import *
from keras.layers import *

def segnet(n_classes, input_height=64, input_width=64):
	img_input = Input(shape=(3, input_height, input_width))
	#Encoder
	x = Conv2D(64, (3, 3), activation='relu', padding='same', name='block1_conv1', data_format='channels_first')(img_input)
	x = (BatchNormalization())(x)
	x = MaxPooling2D((2, 2), strides=(2, 2), name='block1_pool', data_format='channels_first')(x)

	x = Conv2D(128, (3, 3), activation='relu', padding='same', name='block2_conv1', data_format='channels_first')(x)
	x = (BatchNormalization())(x)
	x = MaxPooling2D((2, 2), strides=(2, 2), name='block2_pool', data_format='channels_first')(x)

	x = Conv2D(256, (3, 3), activation='relu', padding='same', name='block3_conv1', data_format='channels_first')(x)
	x = (BatchNormalization())(x)
	x = MaxPooling2D((2, 2), strides=(2, 2), name='block3_pool', data_format='channels_first')(x)

	x = Conv2D(512, (3, 3), activation='relu', padding='same', name='block4_conv1', data_format='channels_first')(x)
	x = (BatchNormalization())(x)
	x = MaxPooling2D((2, 2), strides=(2, 2), name='block4_pool', data_format='channels_first')(x)

	#decoder
	o = Conv2D(512, (3, 3), activation='relu', padding='same', data_format='channels_first')(x)
	o = (BatchNormalization())(o)

	o = (UpSampling2D((2, 2), data_format='channels_first'))(o)
	o = (ZeroPadding2D((1, 1), data_format='channels_first'))(o)
	o = (Conv2D(256, (3, 3), padding='valid', data_format='channels_first'))(o)
	o = (BatchNormalization())(o)

	o = (UpSampling2D((2, 2), data_format='channels_first'))(o)
	o = (ZeroPadding2D((1, 1), data_format='channels_first'))(o)
	o = (Conv2D(128, (3, 3), padding='valid', data_format='channels_first'))(o)
	o = (BatchNormalization())(o)

	o = (UpSampling2D((2, 2), data_format='channels_first'))(o)
	o = (ZeroPadding2D((1, 1), data_format='channels_first'))(o)
	o = (Conv2D(64, (3, 3), padding='valid', data_format='channels_first'))(o)
	o = (BatchNormalization())(o)

	o = Conv2D(n_classes, (3, 3), padding='same', data_format='channels_first')(o)
	o_shape = Model(img_input, o).output_shape
	outputHeight = o_shape[2]
	outputWidth = o_shape[3]
	print(outputHeight)
	print(outputWidth)
	# o = (Reshape((-1, outputHeight * outputWidth)))(o)
	# o = (Permute((2, 1)))(o)
	# o = (Activation('softmax'))(o)
	model = Model(img_input, o)
	model.outputWidth = outputWidth
	model.outputHeight = outputHeight
	
	return model

